package com.huawei.codearts.tools.ratelimit.service;

import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class TokenBucketLimiter {

    private final int maxTokens; // 桶的容量
    private final int refillRate; // 令牌补充速率（每秒补充多少个令牌）
    private AtomicInteger tokens; // 当前令牌数量
    // 定时任务线程池，用于按时补充令牌
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

    public TokenBucketLimiter(int maxTokens, int refillRate) {
        this.maxTokens = maxTokens;
        this.refillRate = refillRate;
        this.tokens = new AtomicInteger(maxTokens);

        // 定时器，每秒执行一次，补充令牌
        scheduler.scheduleAtFixedRate(this::refill, 0, 1, TimeUnit.SECONDS);
    }

    // 尝试获取令牌
    public boolean tryAcquire() {
        int currentTokens = tokens.get();
        if (currentTokens > 0 && tokens.compareAndSet(currentTokens, currentTokens - 1)) {
            return true;
        }
        return false;
    }

    // 补充令牌
    private void refill() {
        int currentTokens = tokens.get();
        int newTokens = Math.min(currentTokens + refillRate, maxTokens);
        tokens.set(newTokens);
    }
}
